"""
这个文件用来读取train_acc.txt，然后画acc和loss曲线
"""
import os

import matplotlib.pyplot as plt
import re

root_path = r'/home/xys/datasets/dataset-larger/CSPdarket/512+B0+BiFPN1+不冻结BN+无数据增强+俩个spp'
log_path = os.path.join(root_path+'/log-yuan.txt')#每一个epoch保存的损失
loss_path = os.path.join(root_path+'/loss-yuan.txt')#训练过程中保存的损失

log_file = open(log_path,'r')
loss_file = open(loss_path,'r')

steps_list, cls_list, cnt_list, reg_list, total_list = [], [], [], [], []
for line in loss_file.readlines():
    numbers = list(map(float, re.findall(r'\d+\.?\d*', line)))
    steps = int(numbers[0])
    cls = numbers[4]
    cnt = numbers[5]
    reg = numbers[6]
    total = numbers[9]

    steps_list.append(steps)
    cls_list.append(cls)
    cnt_list.append(cnt)
    reg_list.append(reg)
    total_list.append(total)
loss_file.close()
epoch_list, traloss_list, valloss_list = [], [], []
for line in log_file.readlines():
    if line == '\n':
        continue
    numbers = list(map(float, re.findall(r'\d+\.?\d*', line)))
    epoch = int(numbers[0])
    traloss = numbers[1]
    valloss = numbers[3]

    epoch_list.append(epoch)
    traloss_list.append(traloss)
    valloss_list.append(valloss)
log_file.close()

#------------------------画图--------------------------------------
fig, axs = plt.subplots(nrows=3, ncols=2, figsize=(15, 4), dpi=100)
# axs[0].plot(epochs_list,lr_list,c='green')
# axs[0].set_title('lr')
axs[0][0].plot(steps_list, cls_list, c='red', label='class_loss')
axs[0][0].set_title('class_loss')
axs[0][1].plot(steps_list, cnt_list, c='green', label='cnt loss')
axs[0][1].set_title('cnt loss')
axs[1][0].plot(steps_list, reg_list, c='blue', label='reg_loss')
axs[1][0].set_title('reg_loss')
axs[1][1].plot(steps_list, total_list, c='purple', label='total loss')
axs[1][1].set_title('total loss')

axs[2][0].plot(epoch_list, traloss_list, c='cyan', label='train loss')
axs[2][0].set_title('train loss')
axs[2][1].plot(epoch_list, valloss_list, c='pink', label='val loss')
axs[2][1].set_title('val loss')
plt.show()
